package middleware

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"happy-book/pkg/ratelimit"
	"net/http"
)

const (
	rateLimitKetFormat = "rate:limit:%s" //限流Key
)

// RateLimitMiddleware 限流middleware
type RateLimitMiddleware struct {
	//限流器组件
	limiter ratelimit.Limiter
}

// NewRateLimitMiddleware 构造函数
func NewRateLimitMiddleware(limiter ratelimit.Limiter) *RateLimitMiddleware {
	return &RateLimitMiddleware{
		limiter: limiter,
	}
}

// RateLimitByIP 根据IP进行限流
func (r *RateLimitMiddleware) RateLimitByIP() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		ip := ctx.ClientIP()
		limit, err := r.limiter.Limit(ctx.Request.Context(), fmt.Sprintf(rateLimitKetFormat, ip))
		if err != nil {
			//TODO 异常埋点
			fmt.Printf("rate limit error: %v\n", err)
			ctx.AbortWithStatus(http.StatusTooManyRequests)
			return
		}
		if limit { //被限流
			//TODO 异常埋点
			fmt.Printf("trigger rate limit! ip: %s\n", ip)
			ctx.AbortWithStatus(http.StatusTooManyRequests)
			return
		}

		//正常放行
		ctx.Next()
	}
}
